import platform
import socket
import subprocess
import psutil
import time
from datetime import datetime
import os
import shutil


def get_linux_distro():
    """
    通过读取 /etc/os-release 文件获取更准确的 Linux 发行版信息。
    """
    try:
        with open('/etc/os-release') as f:
            for line in f:
                if line.startswith('PRETTY_NAME='):
                    # 去掉引号和换行符
                    return line.split('=')[1].strip().strip('"')
    except FileNotFoundError:
        # 如果 /etc/os-release 不存在，则使用旧方法作为备选
        try:
            # hasattr 用于兼容旧版 Python
            if hasattr(platform, 'linux_distribution'):
                distro = platform.linux_distribution()[0]
                if distro:
                    return distro
        except:
            pass  # 忽略旧方法的异常
    return "未知发行版"


def security_scan():
    """执行安全漏洞扫描并显示结果"""
    clear_screen()
    print("=" * 50)
    print("安全漏洞扫描")
    print("=" * 50)

    vulnerabilities = []
    stats = {}

    scan_steps = [
        ("检查系统信息...", scan_system_info),
        ("扫描开放端口...", scan_open_ports),
        ("检查防火墙状态...", scan_firewall_status),
        ("检查系统更新...", scan_system_updates),
        ("扫描运行服务...", scan_running_services),
        ("检查用户账户...", scan_user_accounts)
    ]

    total_steps = len(scan_steps)
    for i, (step_name, scan_func) in enumerate(scan_steps):
        print(f"[{i + 1}/{total_steps}] {step_name}")
        vulnerabilities.extend(scan_func())
        time.sleep(0.3)

    stats = calculate_stats(vulnerabilities)
    display_results(vulnerabilities, stats)
    export_report(vulnerabilities, stats)
    input("\n按 Enter 返回主菜单...")


def clear_screen():
    """清屏，兼容Linux"""
    os.system('clear')


def calculate_stats(vulnerabilities):
    """计算统计数据"""
    risk_counts = {"高": 0, "中": 0, "低": 0}
    type_counts = {}

    for vuln in vulnerabilities:
        risk = vuln.get('risk', '低')
        vuln_type = vuln.get('type', '其他')

        if risk in risk_counts:
            risk_counts[risk] += 1

        if vuln_type in type_counts:
            type_counts[vuln_type] += 1
        else:
            type_counts[vuln_type] = 1

    return {
        'total': len(vulnerabilities),
        'risk_counts': risk_counts,
        'type_counts': type_counts
    }


def scan_system_info():
    """检查系统信息（兼容性增强）"""
    vulnerabilities = []
    try:
        release = platform.release()
        distro = get_linux_distro()

        # 示例：检查是否为旧版本（以 CentOS 7 为例）
        is_old = False
        if 'centos' in distro.lower() and release.startswith('3.10'):  # CentOS 7 内核通常是 3.10
            is_old = True
            description = f'操作系统 {distro} {release} 版本过旧'
            suggestion = '建议升级到更新的系统版本，如 CentOS Stream 或其他现代发行版'
            details = '旧版本可能不再接收关键安全更新'
            risk = '高'
        # 示例：检查是否为旧版本（以Ubuntu为例）
        elif 'ubuntu' in distro.lower() and platform.version().split('.')[0] < "20":
            is_old = True
            description = f'操作系统 {distro} {release} 版本过旧'
            suggestion = '升级到最新的LTS版本'
            details = '旧版本不再接收安全更新'
            risk = '高'

        if is_old:
            vulnerabilities.append({
                'type': '系统版本', 'description': description,
                'risk': risk, 'suggestion': suggestion, 'details': details
            })
        else:
            vulnerabilities.append({
                'type': '系统版本',
                'description': f'操作系统 {distro} {release}',
                'risk': '低',
                'suggestion': '保持系统更新',
                'details': '当前系统版本正常'
            })

        python_version = platform.python_version()
        if tuple(map(int, python_version.split('.'))) < (3, 8, 0):
            vulnerabilities.append({
                'type': '软件版本',
                'description': f'Python {python_version} 版本过旧',
                'risk': '中',
                'suggestion': '升级Python到3.8或更高版本',
                'details': '旧版本可能存在已知的安全漏洞'
            })

    except Exception as e:
        vulnerabilities.append({
            'type': '系统检查', 'description': f'系统信息检查失败: {str(e)}',
            'risk': '中', 'suggestion': '手动检查系统信息',
            'details': '无法自动获取系统信息'
        })
    return vulnerabilities


def scan_open_ports():
    """扫描开放端口"""
    vulnerabilities = []
    # 21(FTP), 23(Telnet) 是非常不安全的协议
    dangerous_ports = {21: 'FTP', 23: 'Telnet', 139: 'NetBIOS', 445: 'SMB'}
    # 常见但需注意的服务端口
    common_ports = {22: 'SSH', 25: 'SMTP', 53: 'DNS', 80: 'HTTP', 110: 'POP3', 143: 'IMAP', 443: 'HTTPS', 3306: 'MySQL',
                    5432: 'PostgreSQL', 3389: 'RDP'}

    ports_to_scan = {**dangerous_ports, **common_ports}

    for port, service_name in ports_to_scan.items():
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
                sock.settimeout(0.5)
                if sock.connect_ex(('127.0.0.1', port)) == 0:
                    risk = '高' if port in dangerous_ports else '中'
                    vulnerabilities.append({
                        'type': '开放端口',
                        'description': f'端口 {port} ({service_name}) 开放',
                        'risk': risk,
                        'suggestion': '确认此端口是否必要开放，如非必要请关闭，或使用防火墙加固访问控制',
                        'details': f'开放的 {service_name} 端口是潜在的攻击入口'
                    })
        except Exception:
            pass  # 忽略扫描过程中的socket错误
    return vulnerabilities


def scan_firewall_status():
    """检查防火墙状态（兼容 ufw 和 firewalld）"""
    vulnerabilities = []
    try:
        # 优先检查 firewalld (CentOS/RHEL/Fedora)
        if shutil.which('firewall-cmd'):
            result = subprocess.run(['firewall-cmd', '--state'], capture_output=True, text=True, encoding='utf-8',
                                    errors='replace')
            if 'running' in result.stdout:
                vulnerabilities.append({
                    'type': '防火墙', 'description': 'Firewalld 防火墙已启用', 'risk': '低',
                    'suggestion': '保持防火墙开启并定期检查规则', 'details': '防火墙状态正常'
                })
            else:
                vulnerabilities.append({
                    'type': '防火墙', 'description': 'Firewalld 防火墙未启用', 'risk': '高',
                    'suggestion': '强烈建议启用防火墙 (sudo systemctl start firewalld && sudo systemctl enable firewalld)',
                    'details': '防火墙是抵御网络攻击的第一道防线'
                })
            return vulnerabilities

        # 其次检查 ufw (Debian/Ubuntu)
        if shutil.which('ufw'):
            result = subprocess.run(['ufw', 'status'], capture_output=True, text=True, encoding='utf-8',
                                    errors='replace')
            if 'Status: active' in result.stdout:
                vulnerabilities.append({
                    'type': '防火墙', 'description': 'UFW 防火墙已启用', 'risk': '低',
                    'suggestion': '保持防火墙开启并定期检查规则', 'details': '防火墙状态正常'
                })
            else:  # 'Status: inactive'
                vulnerabilities.append({
                    'type': '防火墙', 'description': 'UFW 防火墙未启用', 'risk': '高',
                    'suggestion': '启用UFW防火墙 (sudo ufw enable)',
                    'details': '防火墙是抵御网络攻击的第一道防线'
                })
            return vulnerabilities

        # 如果都找不到
        raise FileNotFoundError("未找到 ufw 或 firewall-cmd 命令")

    except Exception as e:
        vulnerabilities.append({
            'type': '防火墙', 'description': f'防火墙状态检查失败: {str(e)}', 'risk': '中',
            'suggestion': '请手动检查防火墙状态 (如: sudo systemctl status firewalld 或 sudo ufw status)',
            'details': '无法自动检测防火墙状态，可能未安装或使用了其他防火墙工具'
        })
    return vulnerabilities


def scan_system_updates():
    """检查系统更新（兼容 apt 和 yum/dnf，并优化 yum/dnf 的退出码处理）"""
    vulnerabilities = []
    try:
        # 优先检查 yum/dnf (CentOS/RHEL/Fedora)
        pkg_manager = None
        if shutil.which('dnf'):
            pkg_manager = 'dnf'
        elif shutil.which('yum'):
            pkg_manager = 'yum'

        if pkg_manager:
            result = subprocess.run([pkg_manager, 'check-update', '-q'], capture_output=True, text=True,
                                    encoding='utf-8', errors='replace')

            # 退出码 100: 发现可用更新
            if result.returncode == 100:
                upgradable_count = len(
                    [line for line in result.stdout.strip().split('\n') if line and not line.isspace()])
                vulnerabilities.append({
                    'type': '系统更新', 'description': f'发现 {upgradable_count} 个待安装的系统更新', 'risk': '中',
                    'suggestion': f'运行 sudo {pkg_manager} update 安装更新',
                    'details': '及时更新有助于修复已知的安全漏洞'
                })
            # 退出码 0: 无可用更新
            elif result.returncode == 0:
                vulnerabilities.append({
                    'type': '系统更新', 'description': '系统更新状态正常', 'risk': '低',
                    'suggestion': '定期检查系统更新', 'details': '当前无待安装的更新'
                })
            # 其他非零退出码: 检查过程出错
            else:
                vulnerabilities.append({
                    'type': '系统更新', 'description': f'更新检查命令执行时发生错误 (退出码: {result.returncode})',
                    'risk': '中',
                    'suggestion': f'请手动运行 "sudo {pkg_manager} check-update" 查看详细错误信息',
                    'details': f'错误可能由网络问题或软件源配置引起。详细输出: {result.stderr or result.stdout}'
                })
            return vulnerabilities

        # 其次检查 apt (Debian/Ubuntu) - 这部分逻辑保持不变
        if shutil.which('apt'):
            result = subprocess.run(['apt', 'list', '--upgradable'], capture_output=True, text=True, encoding='utf-8',
                                    errors='replace')
            upgradable_count = len(
                [line for line in result.stdout.split('\n') if line.strip() and not line.startswith('Listing')])

            if upgradable_count > 0:
                vulnerabilities.append({
                    'type': '系统更新', 'description': f'发现 {upgradable_count} 个待安装的系统更新', 'risk': '中',
                    'suggestion': '运行 sudo apt update && sudo apt upgrade 安装更新',
                    'details': '及时更新有助于修复安全漏洞'
                })
            else:
                vulnerabilities.append({
                    'type': '系统更新', 'description': '系统更新状态正常', 'risk': '低',
                    'suggestion': '定期运行 sudo apt update 检查系统更新', 'details': '当前无待安装的更新'
                })
            return vulnerabilities

        raise FileNotFoundError("未找到 apt, yum, 或 dnf 包管理器")

    except Exception as e:
        vulnerabilities.append({
            'type': '系统更新', 'description': f'更新状态检查失败: {str(e)}', 'risk': '中',
            'suggestion': '请手动检查系统更新 (如: sudo yum update 或 sudo apt update)',
            'details': '无法自动检测更新状态，请确保网络连接正常并有相应权限'
        })
    return vulnerabilities


def scan_running_services():
    """扫描运行中的不安全服务"""
    vulnerabilities = []
    # telnetd, ftpd 等服务进程名可能不同，这里使用通用名称
    dangerous_services = ['telnet', 'ftp', 'rsh', 'rlogin']
    found_services = set()

    try:
        for proc in psutil.process_iter(['name']):
            try:
                proc_name = proc.info['name'].lower() if proc.info['name'] else ''
                for svc in dangerous_services:
                    if svc in proc_name and svc not in found_services:
                        vulnerabilities.append({
                            'type': '危险服务',
                            'description': f'发现可能不安全的服务正在运行: {proc.info["name"]}',
                            'risk': '高',
                            'suggestion': f'若非必要，请停止并禁用此服务 (例如: sudo systemctl stop {svc}d && sudo systemctl disable {svc}d)',
                            'details': f'{svc.upper()} 以明文传输数据，存在严重安全风险'
                        })
                        found_services.add(svc)  # 避免重复报告
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue

        if not vulnerabilities:
            vulnerabilities.append({
                'type': '服务扫描', 'description': '未发现常见的危险服务', 'risk': '低',
                'suggestion': '保持良好习惯，定期审查运行中的服务', 'details': '服务状态正常'
            })


    except Exception as e:
        vulnerabilities.append({
            'type': '服务扫描', 'description': f'服务扫描失败: {str(e)}', 'risk': '低',
            'suggestion': '手动检查运行中的服务 (systemctl list-units --type=service)',
            'details': '建议定期审查系统服务'
        })
    return vulnerabilities


def scan_user_accounts():
    """检查用户账户安全"""
    vulnerabilities = []
    try:
        # 检查空密码用户
        with open('/etc/shadow', 'r', encoding='utf-8') as f:
            for line in f:
                parts = line.strip().split(':')
                if len(parts) > 1:
                    username, password_hash = parts[0], parts[1]
                    # '!' 或 '!!' 表示密码被锁定, '*' 表示非登录用户, 空表示无密码
                    if not password_hash and username != 'root':  # root账户特殊处理
                        vulnerabilities.append({
                            'type': '用户账户', 'description': f'用户 {username} 可能没有设置密码',
                            'risk': '高', 'suggestion': f'为用户设置一个强密码 (sudo passwd {username})',
                            'details': '空密码账户是严重的安全漏洞'
                        })

        # 检查拥有 sudo 权限的用户 (更可靠的方式)
        admin_users = set()
        try:
            # 检查 sudo 组
            result = subprocess.run(['getent', 'group', 'sudo'], capture_output=True, text=True, check=True)
            if result.stdout.strip():
                members = result.stdout.strip().split(':')[-1]
                if members:
                    admin_users.update(members.split(','))
        except (FileNotFoundError, subprocess.CalledProcessError):
            pass  # sudo 组可能不存在

        try:
            # 检查 wheel 组 (CentOS/RHEL)
            result = subprocess.run(['getent', 'group', 'wheel'], capture_output=True, text=True, check=True)
            if result.stdout.strip():
                members = result.stdout.strip().split(':')[-1]
                if members:
                    admin_users.update(members.split(','))
        except (FileNotFoundError, subprocess.CalledProcessError):
            pass  # wheel 组可能不存在

        admin_count = len(admin_users)
        if admin_count > 3:  # 阈值可调整
            vulnerabilities.append({
                'type': '用户账户', 'description': f'发现 {admin_count} 个特权用户', 'risk': '中',
                'suggestion': '遵循最小权限原则，移除不必要的特权用户 (sudo gpasswd -d <username> <group>)',
                'details': f'过多的特权用户增加了系统被攻破的风险。特权用户列表: {", ".join(admin_users)}'
            })

    except PermissionError:
        vulnerabilities.append({
            'type': '用户账户', 'description': '用户账户检查失败: 权限不足',
            'risk': '中', 'suggestion': '请使用 root 或 sudo 权限运行此脚本以检查 /etc/shadow 文件',
            'details': '没有权限读取 /etc/shadow 文件，无法检查空密码用户'
        })
    except Exception as e:
        vulnerabilities.append({
            'type': '用户账户', 'description': f'用户账户检查失败: {str(e)}', 'risk': '低',
            'suggestion': '手动检查用户账户 (cat /etc/passwd 和 getent group sudo)',
            'details': '建议定期审查用户权限'
        })
    return vulnerabilities


def display_results(vulnerabilities, stats):
    """显示扫描结果"""
    clear_screen()
    print("=" * 50)
    print("安全漏洞扫描结果")
    print("=" * 50)
    print(f"扫描时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"总问题数: {stats['total']}")
    print(f"  - 高风险: {stats['risk_counts']['高']}")
    print(f"  - 中风险: {stats['risk_counts']['中']}")
    print(f"  - 低风险: {stats['risk_counts']['低']}")
    print("\n" + "=" * 50)
    print("详细问题列表:")
    print("-" * 50)

    if not vulnerabilities:
        print("恭喜！未发现明显的安全问题。")
        print("-" * 50)
        return

    # 按风险等级排序
    vulnerabilities.sort(key=lambda x: {"高": 0, "中": 1, "低": 2}.get(x.get('risk', '低'), 3))

    for i, vuln in enumerate(vulnerabilities, 1):
        print(f"{i}. [{vuln.get('type', '未知类型')}] {vuln.get('description', '无描述')}")
        print(f"   风险等级: {vuln.get('risk', '未知')}")
        print(f"   建议措施: {vuln.get('suggestion', '无')}")
        print(f"   详细信息: {vuln.get('details', '无')}")
        print("-" * 50)


def export_report(vulnerabilities, stats):
    """导出扫描报告到文本文件"""
    try:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"security_scan_report_{timestamp}.txt"

        with open(filename, 'w', encoding='utf-8') as f:
            f.write("=" * 50 + "\n")
            f.write("安全漏洞扫描报告\n")
            f.write("=" * 50 + "\n\n")
            f.write(f"扫描时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
            f.write("-" * 20 + "\n")
            f.write("扫描结果统计:\n")
            f.write(f"  总问题数: {stats['total']}\n")
            f.write(f"  高风险: {stats['risk_counts']['高']}\n")
            f.write(f"  中风险: {stats['risk_counts']['中']}\n")
            f.write(f"  低风险: {stats['risk_counts']['低']}\n\n")
            f.write("=" * 50 + "\n")
            f.write("详细问题列表:\n")
            f.write("-" * 50 + "\n\n")

            if not vulnerabilities:
                f.write("恭喜！未发现明显的安全问题。\n")
            else:
                # 报告中也按风险排序
                vulnerabilities.sort(key=lambda x: {"高": 0, "中": 1, "低": 2}.get(x.get('risk', '低'), 3))
                for i, vuln in enumerate(vulnerabilities, 1):
                    f.write(f"{i}. [{vuln.get('type', '未知类型')}] {vuln.get('description', '无描述')}\n")
                    f.write(f"   风险等级: {vuln.get('risk', '未知')}\n")
                    f.write(f"   建议措施: {vuln.get('suggestion', '无')}\n")
                    f.write(f"   详细信息: {vuln.get('details', '无')}\n")
                    f.write("-" * 50 + "\n")

        print(f"\n报告已成功保存为当前目录下的文件: {filename}")

    except Exception as e:
        print(f"\n导出报告失败: {str(e)}")